package com.geo.ai.adapter.impl;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.geo.ai.adapter.AIModelAdapter;
import com.geo.ai.adapter.ModelUsageStats;
import com.geo.ai.adapter.StreamCallback;
import com.geo.ai.model.AIRequest;
import com.geo.ai.model.AIResponse;
import com.geo.ai.model.ModelConfig;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.security.MessageDigest;
import java.time.Duration;
import java.time.LocalDateTime;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 百度文心一言模型适配器实现
 * 
 * @author GEO Team
 * @version 1.0
 */
@Slf4j
@Component
public class WenxinAdapter implements AIModelAdapter {
    
    private static final org.slf4j.Logger log = org.slf4j.LoggerFactory.getLogger(WenxinAdapter.class);
    
    private final OkHttpClient httpClient;
    private final ObjectMapper objectMapper;
    private final ModelConfig modelConfig;
    private final String secretKey;
    
    // 访问令牌缓存
    private String accessToken;
    private long tokenExpireTime;
    
    // 使用统计
    private final AtomicLong totalRequests = new AtomicLong(0);
    private final AtomicLong successfulRequests = new AtomicLong(0);
    private final AtomicLong failedRequests = new AtomicLong(0);
    private final AtomicLong totalTokensUsed = new AtomicLong(0);
    private final AtomicLong totalResponseTime = new AtomicLong(0);
    private double totalCost = 0.0;
    
    public WenxinAdapter(@Value("${ai.models.wenxin.api-key}") String apiKey,
                        @Value("${ai.models.wenxin.secret-key}") String secretKey,
                        @Value("${ai.models.wenxin.base-url}") String baseUrl,
                        @Value("${ai.models.wenxin.model}") String model,
                        @Value("${ai.models.wenxin.enabled:true}") Boolean enabled,
                        @Value("${ai.models.wenxin.priority:3}") Integer priority,
                        @Value("${ai.models.wenxin.cost-per-token:0.008}") Double costPerToken,
                        @Value("${ai.models.wenxin.rate-limit-rpm:60}") Integer rateLimitRpm) {
        
        this.httpClient = new OkHttpClient.Builder()
                .connectTimeout(Duration.ofSeconds(30))
                .readTimeout(Duration.ofSeconds(120))
                .writeTimeout(Duration.ofSeconds(30))
                .build();
        
        this.objectMapper = new ObjectMapper();
        this.secretKey = secretKey;
        
        this.modelConfig = ModelConfig.builder()
                .modelName(model)
                .provider("Baidu")
                .apiEndpoint(baseUrl)
                .apiKey(apiKey)
                .enabled(enabled)
                .priority(priority)
                .costPerToken(costPerToken)
                .rateLimitRpm(rateLimitRpm)
                .maxTokens(2048)
                .defaultTemperature(0.7)
                .timeoutMs(120000L)
                .retryTimes(3)
                .capabilities(ModelConfig.ModelCapabilities.builder()
                        .supportsStreaming(true)
                        .supportsFunctionCalling(false)
                        .supportsImageInput(false)
                        .supportsMultimodal(false)
                        .contextWindow(2048)
                        .maxOutputLength(2048)
                        .build())
                .build();
    }
    
    @Override
    public String getAdapterName() {
        return "Wenxin-ERNIE";
    }
    
    @Override
    public String getModelType() {
        return "text-generation";
    }
    
    @Override
    public boolean isAvailable() {
        return modelConfig.isValid() && validateConfig();
    }
    
    @Override
    public ModelConfig getModelConfig() {
        return modelConfig;
    }
    
    @Override
    public AIResponse invoke(AIRequest request) {
        long startTime = System.currentTimeMillis();
        totalRequests.incrementAndGet();
        
        try {
            // 获取访问令牌
            String token = getAccessToken();
            if (token == null) {
                failedRequests.incrementAndGet();
                return buildErrorResponse(request, "AUTH_ERROR", "无法获取访问令牌");
            }
            
            // 构建请求JSON
            String requestJson = buildRequestJson(request);
            
            // 创建HTTP请求
            String apiUrl = modelConfig.getApiEndpoint() + "/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" + token;
            Request httpRequest = new Request.Builder()
                    .url(apiUrl)
                    .post(RequestBody.create(requestJson, MediaType.parse("application/json")))
                    .addHeader("Content-Type", "application/json")
                    .build();
            
            // 执行请求
            try (Response response = httpClient.newCall(httpRequest).execute()) {
                long responseTime = System.currentTimeMillis() - startTime;
                totalResponseTime.addAndGet(responseTime);
                
                if (!response.isSuccessful()) {
                    failedRequests.incrementAndGet();
                    return buildErrorResponse(request, "HTTP_ERROR", "HTTP Error: " + response.code());
                }
                
                String responseBody = response.body().string();
                return parseResponse(request, responseBody, responseTime);
            }
            
        } catch (Exception e) {
            failedRequests.incrementAndGet();
            log.error("文心一言API调用失败", e);
            return buildErrorResponse(request, "API_ERROR", e.getMessage());
        }
    }
    
    @Override
    public CompletableFuture<AIResponse> invokeAsync(AIRequest request) {
        return CompletableFuture.supplyAsync(() -> invoke(request));
    }
    
    @Override
    public void invokeStream(AIRequest request, StreamCallback callback) {
        try {
            String token = getAccessToken();
            if (token == null) {
                callback.onError(new RuntimeException("无法获取访问令牌"));
                return;
            }
            
            String requestJson = buildStreamRequestJson(request);
            String apiUrl = modelConfig.getApiEndpoint() + "/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" + token;
            
            Request httpRequest = new Request.Builder()
                    .url(apiUrl)
                    .post(RequestBody.create(requestJson, MediaType.parse("application/json")))
                    .addHeader("Content-Type", "application/json")
                    .build();
            
            httpClient.newCall(httpRequest).enqueue(new Callback() {
                @Override
                public void onFailure(Call call, IOException e) {
                    callback.onError(e);
                }
                
                @Override
                public void onResponse(Call call, Response response) throws IOException {
                    if (!response.isSuccessful()) {
                        callback.onError(new RuntimeException("HTTP Error: " + response.code()));
                        return;
                    }
                    
                    String fullResponse = response.body().string();
                    callback.onChunk(fullResponse);
                    
                    AIResponse aiResponse = parseResponse(request, fullResponse, 0L);
                    callback.onComplete(aiResponse);
                }
            });
            
        } catch (Exception e) {
            callback.onError(e);
        }
    }
    
    @Override
    public boolean validateConfig() {
        try {
            String token = getAccessToken();
            return token != null && !token.isEmpty();
        } catch (Exception e) {
            log.warn("文心一言配置验证失败", e);
            return false;
        }
    }
    
    @Override
    public ModelUsageStats getUsageStats() {
        long avgResponseTime = totalRequests.get() > 0 ? 
                totalResponseTime.get() / totalRequests.get() : 0;
        
        return new ModelUsageStats(
                totalRequests.get(),
                successfulRequests.get(),
                failedRequests.get(),
                totalTokensUsed.get(),
                totalCost,
                avgResponseTime
        );
    }
    
    @Override
    public void resetUsageStats() {
        totalRequests.set(0);
        successfulRequests.set(0);
        failedRequests.set(0);
        totalTokensUsed.set(0);
        totalResponseTime.set(0);
        totalCost = 0.0;
    }
    
    /**
     * 获取访问令牌
     */
    private String getAccessToken() {
        // 检查token是否过期
        if (accessToken != null && System.currentTimeMillis() < tokenExpireTime) {
            return accessToken;
        }
        
        try {
            String tokenUrl = "https://aip.baidubce.com/oauth/2.0/token?" +
                    "grant_type=client_credentials&" +
                    "client_id=" + modelConfig.getApiKey() + "&" +
                    "client_secret=" + secretKey;
            
            Request request = new Request.Builder()
                    .url(tokenUrl)
                    .post(RequestBody.create("", MediaType.parse("application/json")))
                    .build();
            
            try (Response response = httpClient.newCall(request).execute()) {
                if (response.isSuccessful()) {
                    String responseBody = response.body().string();
                    JsonNode jsonResponse = objectMapper.readTree(responseBody);
                    
                    accessToken = jsonResponse.get("access_token").asText();
                    int expiresIn = jsonResponse.get("expires_in").asInt();
                    
                    // 提前5分钟过期，避免边界情况
                    tokenExpireTime = System.currentTimeMillis() + (expiresIn - 300) * 1000L;
                    
                    return accessToken;
                }
            }
        } catch (Exception e) {
            log.error("获取文心一言访问令牌失败", e);
        }
        
        return null;
    }
    
    /**
     * 构建请求JSON
     */
    private String buildRequestJson(AIRequest request) throws Exception {
        var jsonObject = objectMapper.createObjectNode();
        
        // 文心一言使用不同的参数结构
        if (request.getMessages() != null && !request.getMessages().isEmpty()) {
            // 获取最后一条用户消息
            AIRequest.Message lastUserMessage = request.getMessages().stream()
                    .filter(msg -> msg.getRole() == AIRequest.Message.Role.USER)
                    .reduce((first, second) -> second)
                    .orElse(null);
            
            if (lastUserMessage != null) {
                jsonObject.put("message", lastUserMessage.getContent());
            }
        } else if (request.getUserInput() != null) {
            jsonObject.put("message", request.getUserInput());
        }
        
        jsonObject.put("temperature", request.getTemperature() != null ? request.getTemperature() : modelConfig.getDefaultTemperature());
        jsonObject.put("top_p", 0.8);
        jsonObject.put("penalty_score", 1.0);
        jsonObject.put("stream", false);
        
        return objectMapper.writeValueAsString(jsonObject);
    }
    
    /**
     * 构建流式请求JSON
     */
    private String buildStreamRequestJson(AIRequest request) throws Exception {
        var jsonObject = objectMapper.createObjectNode();
        
        if (request.getMessages() != null && !request.getMessages().isEmpty()) {
            AIRequest.Message lastUserMessage = request.getMessages().stream()
                    .filter(msg -> msg.getRole() == AIRequest.Message.Role.USER)
                    .reduce((first, second) -> second)
                    .orElse(null);
            
            if (lastUserMessage != null) {
                jsonObject.put("message", lastUserMessage.getContent());
            }
        } else if (request.getUserInput() != null) {
            jsonObject.put("message", request.getUserInput());
        }
        
        jsonObject.put("temperature", request.getTemperature() != null ? request.getTemperature() : modelConfig.getDefaultTemperature());
        jsonObject.put("top_p", 0.8);
        jsonObject.put("penalty_score", 1.0);
        jsonObject.put("stream", true);
        
        return objectMapper.writeValueAsString(jsonObject);
    }
    
    /**
     * 解析响应
     */
    private AIResponse parseResponse(AIRequest request, String responseBody, long responseTime) {
        try {
            JsonNode jsonResponse = objectMapper.readTree(responseBody);
            
            // 检查是否有错误
            if (jsonResponse.has("error_code")) {
                String errorMsg = jsonResponse.get("error_msg").asText();
                return buildErrorResponse(request, "API_ERROR", errorMsg);
            }
            
            String content = jsonResponse.get("result").asText();
            
            // 解析token使用情况
            AIResponse.TokenUsage tokenUsage = null;
            JsonNode usage = jsonResponse.get("usage");
            if (usage != null) {
                tokenUsage = AIResponse.TokenUsage.builder()
                        .promptTokens(usage.get("prompt_tokens").asInt())
                        .completionTokens(usage.get("completion_tokens").asInt())
                        .totalTokens(usage.get("total_tokens").asInt())
                        .build();
                
                totalTokensUsed.addAndGet(tokenUsage.getTotalTokens());
                totalCost += tokenUsage.getTotalTokens() * modelConfig.getCostPerToken();
            } else {
                // 文心一言可能不返回token信息，进行估算
                int estimatedTokens = content.length() / 2; // 简单估算
                tokenUsage = AIResponse.TokenUsage.builder()
                        .promptTokens(0)
                        .completionTokens(estimatedTokens)
                        .totalTokens(estimatedTokens)
                        .build();
                
                totalTokensUsed.addAndGet(estimatedTokens);
                totalCost += estimatedTokens * modelConfig.getCostPerToken();
            }
            
            successfulRequests.incrementAndGet();
            
            return AIResponse.builder()
                    .requestId(request.getRequestId())
                    .responseId(jsonResponse.has("id") ? jsonResponse.get("id").asText() : null)
                    .modelName(modelConfig.getModelName())
                    .status(AIResponse.ResponseStatus.SUCCESS)
                    .content(content)
                    .tokenUsage(tokenUsage)
                    .responseTimeMs(responseTime)
                    .cost(tokenUsage.getTotalTokens() * modelConfig.getCostPerToken())
                    .responseTime(LocalDateTime.now())
                    .finishReason(AIResponse.FinishReason.STOP)
                    .build();
                    
        } catch (Exception e) {
            log.error("解析文心一言响应失败", e);
            return buildErrorResponse(request, "PARSE_ERROR", e.getMessage());
        }
    }
    
    /**
     * 构建错误响应
     */
    private AIResponse buildErrorResponse(AIRequest request, String errorCode, String errorMessage) {
        return AIResponse.builder()
                .requestId(request.getRequestId())
                .modelName(modelConfig.getModelName())
                .status(AIResponse.ResponseStatus.FAILED)
                .errorCode(errorCode)
                .errorMessage(errorMessage)
                .responseTime(LocalDateTime.now())
                .build();
    }
}